"""
motor MongoDB 二次封装
"""

import math

from motor.core import (
    AgnosticCollection,
    AgnosticCommandCursor,
    AgnosticCursor,
)
from pymongo import (
    DeleteMany,
    DeleteOne,
    InsertOne,
    ReplaceOne,
    UpdateMany,
    UpdateOne,
)
from pymongo.results import (
    BulkWriteResult,
    DeleteResult,
    InsertManyResult,
    InsertOneResult,
    UpdateResult,
)

from db_base.async_mongodb_db.client import async_mongo_db
from my_typing import *


class empty:  # noqa
    """
    该类用于表示没有为给定的输入或输出值提供数据
    它是必需的，因为 `None` 可能是一个有效的输入或输出值
    """

    pass


class CollectionSimpleUtils:
    """ 集合简单的工具类 """

    projection: DictStrInt = {
        "_id": 0,
    }
    """ 字段展示控制（默认不展示 `_id`） """

    deleted_mark: DictStrInt = {
        "is_deleted": 1,
    }
    """ 数据删除标记 """

    undeleted_mark: DictStrInt = {
        "is_deleted": 0,
    }
    """ 数据未删除标记 """

    undeleted_query: dict[str, DictStrInt] = {
        "is_deleted": {
            "$ne": 1,
        },
    }
    """ 数据未删除查询条件 """

    @staticmethod
    def get_time_dict(*args: str) -> DictStrStr:
        """
        获取时间字段的字典

        Args:
            args: 字段

        Returns:
            时间字段的字典

        Examples:
            >>> # 不传任何字段，返回空字典
            >>> CollectionSimpleUtils.get_time_dict()
            {}

            >>> # 传入 `update_time`，返回只有 1 个元素的字典
            >>> CollectionSimpleUtils.get_time_dict("update_time")
            {
                "update_time": "2023-05-10 15:31:23",
            }

            >>> # 传入 `update_time`、`update_time`，返回有 2 个元素的字典
            >>> CollectionSimpleUtils.get_time_dict("create_time", "update_time")
            {
                "create_time": "2023-05-10 15:31:23",
                "update_time": "2023-05-10 15:31:23",
            }
        """

        if not args:
            return {}

        now_datetime_str: str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        """ 当前日期时间（格式: %Y-%m-%d %H:%M:%S） """

        return {
            time_field: now_datetime_str
            for time_field in args
        }

    @staticmethod
    def get_deleted_mark_dict(filter_all: bool) -> dict:
        """
        获取删除标记的字典

        Args:
            filter_all: 是否查询全部数据

        Returns:
            删除标记的字典

        Examples:
            >>> # `filter_all` 为 False，获取未删除的数据，即：`is_deleted` 不等于 1 的数据（兼容没有 `is_deleted` 的字段）
            >>> CollectionSimpleUtils.get_deleted_mark_dict(
            ...     filter_all=False,
            ... )
            {
                "is_deleted": {
                    "$ne": 1,
                },
            }

            >>> # `filter_all` 为 True，获取全部数据
            >>> CollectionSimpleUtils.get_deleted_mark_dict(
            ...     filter_all=True,
            ... )
            {}
        """

        if filter_all:
            return {}

        return CollectionSimpleUtils.undeleted_query.copy()

    @staticmethod
    def _get_filter(
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        filter_all: bool = False,
    ) -> dict:
        """
        获取查询条件

        Args:
            filter: 过滤条件
            filter_all: 是否查询全部数据

        Returns:
            查询条件

        Examples:
            >>> # `filter` 为 None
            >>> CollectionSimpleUtils._get_filter()
            {
                "is_deleted": {
                    "$ne": 1,
                },
            }

            >>> # `filter` 为 None，`filter_all` 为 True
            >>> CollectionSimpleUtils._get_filter(filter_all=True)
            {}

            >>> # `filter` 不为 None，`filter_all` 为 False
            >>> CollectionSimpleUtils._get_filter(
            ...     filter={
            ...         "name": "张三",
            ...         "age": 18,
            ...     },
            ... )
            {
                "name": "张三",
                "age": 18,
                "is_deleted": {
                    "$ne": 1,
                },
            }
        """

        return {
            **(
                {}
                if filter is None
                else filter
            ),
            **CollectionSimpleUtils.get_deleted_mark_dict(filter_all)
        }

    @staticmethod
    def _get_projection(projection: DictOrNone = empty) -> DictOrNone:
        """
        获取展示/不展示的字段字典

        Args:
            projection: 展示/不展示的字段

        Returns:
            展示/不展示的字段字典

        Examples:
            >>> # `projection` 为 empty
            >>> CollectionSimpleUtils._get_projection()
            {
                "_id": 0,
            }

            >>> # `projection` 为 None
            >>> CollectionSimpleUtils._get_projection(projection=None)
            None

            >>> # `projection` 不为 None
            >>> CollectionSimpleUtils._get_projection(
            ...     projection={
            ...         "name": 1,
            ...         "age": 1,
            ...     },
            ... )
            {
                "_id": 0,
                "name": 1,
                "age": 1,
            }

            >>> # `projection` 不为 None，且需要展示 `_id` 字段
            >>> CollectionSimpleUtils._get_projection(
            ...     projection={
            ...         "_id": 1,
            ...         "name": 1,
            ...         "age": 1,
            ...     },
            ... )
            {
                "_id": 1,
                "name": 1,
                "age": 1,
            }
        """

        if projection is empty:
            return CollectionSimpleUtils.projection.copy()

        return projection

    @staticmethod
    def projection_shows(*args: str) -> DictStrInt:
        """
        需要展示的字段（默认不展示 `_id` 字段）

        Args:
            args: 需要展示的字段

        Returns:
            需要展示的字段-字典

        Examples:
            >>> # 需要展示的字段：`name`、`age`
            >>> CollectionSimpleUtils.projection_shows("name", "age")
            {
                "_id": 0,
                "name": 1,
                "age": 1,
            }

            >>> # 需要展示的字段：`_id`、`name`、`age`
            >>> CollectionSimpleUtils.projection_shows("_id", "name", "age")
            {
                "_id": 1,
                "name": 1,
                "age": 1,
            }
        """

        return {
            **CollectionSimpleUtils.projection,
            **{
                field: 1
                for field in args
            },
        }

    @staticmethod
    def projection_no_shows(*args: str) -> DictStrInt:
        """
        不需要展示的字段（默认不展示 _id 字段）

        Args:
            args: 不需要展示的字段

        Returns:
            不需要展示的字段-字典

        Examples:
            >>> # 不需要展示的字段：`name`、`age`
            >>> CollectionSimpleUtils.projection_no_shows("name", "age")
            {
                "_id": 0,
                "name": 0,
                "age": 0,
            }
        """

        return {
            **CollectionSimpleUtils.projection,
            **{
                field: 0
                for field in args
            },
        }


class AsyncCollectionBase(CollectionSimpleUtils):
    """
    异步 MongoDB 集合 基类
    """

    collection_name: str = ""
    """ 集合名称 """

    def __init__(self) -> None:
        self.collection: AgnosticCollection = async_mongo_db[self.collection_name]
        """ 集合 """

    def insert_one(
        self,
        document: dict,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, InsertOneResult]:
        """
        新增 单个 文档

        Args:
            document: 文档

        Returns:
            插入文档的结果（InsertOneResult）

                - acknowledged (bool): 是否成功
                - inserted_id (ObjectId): 插入文档的 _id
        """

        _document: dict = {
            **document,
            **self.undeleted_mark,
            **self.get_time_dict("create_time", "update_time"),
        }

        return self.collection.insert_one(_document, *args, **kwargs)

    def insert_many(
        self,
        documents: Iterable[dict],
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, InsertManyResult]:
        """
        新增 多个 文档

        Args:
            documents: 文档列表

        Returns:
            插入文档的结果（InsertManyResult）

                - `acknowledged` (bool): 是否成功
                - `inserted_ids` (list[ObjectId]): 插入文档的 _id
        """

        _documents: TupleDict = tuple(
            {
                **document,
                **self.undeleted_mark,
                **self.get_time_dict("create_time", "update_time"),
            }
            for document in documents
        )

        return self.collection.insert_many(_documents, *args, **kwargs)

    def update_one(
        self,
        filter: dict,  # noqa: BuiltinShadowName
        update: dict,
        upsert: bool = False,
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, UpdateResult]:
        """
        更新 单个 文档

        Args:
            filter: 过滤条件
            update: 更新的文档
            upsert: 更新插入（存在：更新 / 不存在：创建）
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            更新文档的结果（UpdateResult）

                - `acknowledged` (bool): 是否成功
                - `matched_count` (int): 匹配到的文档的数量
                - `modified_count` (int): 更新文档的数量
                - `upserted_id` (ObjectId): 更新文档的数量
        """

        _filter: dict = self._get_filter(filter, filter_all)
        _update: dict = {
            "$set": {
                **update,
                **self.get_time_dict("update_time"),
            },
        }

        return self.collection.update_one(_filter, _update, upsert, *args, **kwargs)

    def update_many(
        self,
        filter: dict,  # noqa: BuiltinShadowName
        update: dict,
        upsert: bool = False,
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, UpdateResult]:
        """
        更新 多个 文档

        Args:
            filter: 过滤条件
            update: 更新的文档
            upsert: 更新插入（存在：更新 / 不存在：创建）
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            更新文档的结果（UpdateResult）

                - `acknowledged` (bool): 是否成功
                - `matched_count` (int): 更新文档的数量
                - `modified_count` (int): 更新文档的数量
        """

        _filter: dict = self._get_filter(filter, filter_all)
        _update: dict = {
            "$set": {
                **update,
                **self.get_time_dict("update_time"),
            },
        }

        return self.collection.update_many(_filter, _update, upsert, *args, **kwargs)

    def delete_one(
        self,
        filter: dict,  # noqa: BuiltinShadowName
        filter_all: bool = False,
        forced: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, DeleteResult] | Coroutine[Any, Any, UpdateResult]:
        """
        删除 单个 文档（默认软删除）

        Args:
            filter: 过滤条件
            filter_all: 是否查询全部数据（包括已删除的数据）
            forced: 是否强制删除

        Returns:
            删除/更新文档的结果
                DeleteResult: 删除文档的结果
                    - `acknowledged` (bool): 是否成功
                    - `deleted_count` (int): 删除文档的数量

                UpdateResult: 更新文档的结果
                    - `acknowledged` (bool): 是否成功
                    - `matched_count` (int): 匹配到的文档的数量
                    - `modified_count` (int): 更新文档的数量
                    - `upserted_id` (ObjectId): 更新文档的数量
        """

        _filter: dict = self._get_filter(filter, filter_all)

        if forced:
            return self.collection.delete_one(_filter, *args, **kwargs)

        return self.update_one(_filter, self.deleted_mark, filter_all=filter_all)

    def delete_many(
        self,
        filter: dict,  # noqa: BuiltinShadowName
        filter_all: bool = False,
        forced: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, DeleteResult] | Coroutine[Any, Any, UpdateResult]:
        """
        删除 多个 文档（默认软删除）

        Args:
            filter: 过滤条件
            filter_all: 是否查询全部数据（包括已删除的数据）
            forced: 是否强制删除

        Returns:
            删除/更新文档的结果
                DeleteResult: 删除文档的结果
                    - `acknowledged` (bool): 是否成功
                    - `deleted_count` (int): 删除文档的数量

                UpdateResult: 更新文档的结果
                    - `acknowledged` (bool): 是否成功
                    - `matched_count` (int): 匹配到的文档的数量
                    - `modified_count` (int): 更新文档的数量
                    - `upserted_id` (ObjectId): 更新文档的数量
        """

        _filter: dict = CollectionSimpleUtils._get_filter(filter, filter_all)

        if forced:
            return self.collection.delete_many(_filter, *args, **kwargs)

        return self.update_many(_filter, self.deleted_mark, filter_all=filter_all)

    def find_one(
        self,
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        projection: DictStrInt | None = empty,
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, dict]:
        """
        查询 单个 文档（默认不返回 _id）

        Args:
            filter: 过滤条件
            projection: 展示/不展示的字段字典
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            查询到的文档
        """

        _filter: dict = self._get_filter(filter, filter_all)
        _projection: DictStrInt | None = self._get_projection(projection)

        return self.collection.find_one(_filter, projection=_projection, *args, **kwargs) or {}

    async def find(
        self,
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        projection: DictStrInt | None = empty,
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> ListDict:
        """
        查询 多个 文档（默认不返回 _id）

        Args:
            filter: 过滤条件
            projection: 展示/不展示的字段字典
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            查询到的文档
        """

        _filter: dict = self._get_filter(filter, filter_all)
        _projection: DictStrInt | None = self._get_projection(projection)

        _documents_cursor: AgnosticCursor = self.collection.find(_filter, projection=_projection, *args, **kwargs)

        document_list = [
            document
            async for document in _documents_cursor
        ]

        return document_list

    def count_documents(
        self,
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, int]:
        """
        查询 文档总数

        Args:
            filter: 过滤条件
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            文档总数
        """

        _filter: dict = self._get_filter(filter, filter_all)

        return self.collection.count_documents(_filter, *args, **kwargs)

    async def aggregate(
        self,
        pipeline: ListDict,
        *args: Any,
        **kwargs: Any,
    ) -> ListDict:
        """
        聚合查询

        Args:
            pipeline: 聚合管道

        Returns:
            查询到的文档
        """

        _result_cursor: AgnosticCommandCursor = self.collection.aggregate(pipeline, *args, **kwargs)

        result_list = [
            document
            async for document in _result_cursor
        ]

        return result_list

    def distinct(
        self,
        key: str,
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        filter_all: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, list]:
        """
        获取去重后的字段值列表

        Args:
            key: 字段名
            filter: 过滤条件
            filter_all: 是否查询全部数据（包括已删除的数据）

        Returns:
            去重后的字段值列表
        """

        _filter: dict = self._get_filter(filter, filter_all)

        return self.collection.distinct(key, _filter, *args, **kwargs)

    def bulk_write(
        self,
        requests: Sequence[InsertOne | DeleteOne | DeleteMany | ReplaceOne | UpdateOne | UpdateMany],
        *args: Any,
        **kwargs: Any,
    ) -> Coroutine[Any, Any, BulkWriteResult]:
        """
        批量写入

        Args:
            requests: 批量写入请求

        Returns:
            批量写入结果

                - `bulk_api_result` (dict): 批量写入的结果
                - `inserted_count` (int): 插入的文档数
                - `matched_count` (int): 匹配到的文档数
                - `modified_count` (int): 修改的文档数
                - `deleted_count` (int): 删除的文档数
                - `upserted_count` (int): 插入的新文档数
                - `upserted_ids` (DictOrNone): 插入的新文档的 _id 字典
        """

        return self.collection.bulk_write(requests, *args, **kwargs)

    async def paginated_query(
        self,
        filter: DictOrNone = None,  # noqa: BuiltinShadowName
        projection: DictStrInt | None = empty,
        order: list[dict[str, IntOrStr]] | None = None,
        filter_all: bool = False,
        page_index: int = 1,
        page_size: int = 10,
        *args,
        **kwargs,
    ) -> dict[str, ListDict | int]:
        """
        分页查询

        Args:
            filter: 过滤条件
            projection: 展示/不展示的字段字典
            order: 排序字段
            filter_all: 是否查询全部数据（包括已删除的数据）
            page_index: 页码
            page_size: 每页数据条数

        Returns:
            分页后的数据
        """

        _filter: dict = self._get_filter(filter, filter_all)
        _projection: DictStrInt | None = self._get_projection(projection)

        data_count: int = await self.count_documents(_filter)
        """ 总数据条数 """

        page_count: int = math.ceil(data_count / page_size)
        """ 总页数 """

        skip_count: int = page_size * (page_index - 1)
        """ 当前查询需要跳过的数据条数 """

        raw_data: AgnosticCursor = self.collection.find(_filter, _projection, *args, **kwargs)
        """ 查询的原始数据：AgnosticCursor 对象，类似于迭代器，只能循环 1 次 """

        # 排序
        if order is not None:
            raw_data = raw_data.sort(
                key_or_list=[
                    (
                        str(item["field"]),
                        int(item["direction"]),
                    )
                    for item in order
                ]
            )

        data_list: ListDict = [
            i
            async for i in raw_data.skip(skip_count).limit(page_size)
        ]
        """ 分页后的数据 """

        return {
            "page": page_index,  # ->       第几页
            "page_size": page_size,  # ->   每页几条
            "page_count": page_count,  # -> 总页数
            "data_count": data_count,  # -> 总数据条数
            "data_list": data_list,  # ->   分页后的数据
        }
